import sys
import os
import pcapy
import binascii
import datetime
import socket
import struct
import string
from impacket import ImpactDecoder
from hashlib import md5
from time import time
import imp

events = []
prev_capture_time = 0
setting = None

global local_port
local_port = {}
	
class Packet:
	def __init__(self, ip, port, data, len, captured_raw_size, send):
		self.ip = ip
		self.length = len
		self.port = port
		self.data = data
		self.send = send
		self.captured_raw_size = captured_raw_size

class Time:
	def __init__(self, time):
		self.time = time

def get_device():
	# Get the list of devices we can listen on
	ifs = pcapy.findalldevs()

	# No devices found
	if len(ifs) == 0:
		raise RuntimeError, "Error: no available network devices, or you don't have enough permissions on this system."

	# A single device was found
	if len(ifs) == 1:
		device = ifs[0]

	if hasattr(setting, 'device_no'):
		device = ifs[setting.device_no-1]
	else:
		print "No 'device_no' value in 'setting.py'. Choose one!"
		print "Available network devices:"
		for i in xrange(len(ifs)):
			print '\t%i - %s' % (i + 1, ifs[i])
		print
		while 1:
			choice = raw_input("Choose an device [0 to quit]: ")
			try:
				i = int(choice)
				if i == 0:
					device = None
					break
				device = ifs[i-1]
				break
			except Exception:
				pass

	# Return the selected device
	return device

def capture(device):
	decoder = ImpactDecoder.EthDecoder()
	local_ip = socket.gethostbyname(socket.gethostname())

	# Arguments here are:
	#   device
	#   snaplen (maximum number of bytes to capture _per_packet_)
	#   promiscious mode (1 for true)
	#   timeout (in milliseconds)
	promiscuous = False
	read_timeout = 100
	global prev_capture_time
	prev_capture_time = time()
	cap = pcapy.open_live(device, max_cap_length, promiscuous, read_timeout)
	cap.setfilter(filter)
	packet_limit = -1 # infinite
	cap.loop(packet_limit, recv_packets) # capture packets


def append_message(raw_data, raw_size, payload, svr_ip, svr_port, captured_raw_size, to_server):
	#print(rl)
	#print ( "CRC0: ", tup.getCrc1(raw_data ) )
	if to_server:
		if client_decrypt:
			if client_decrypt.run(raw_data, raw_size):
				raw_data = client_decrypt.result()
			else:
				dump_message(payload, raw_size)
				raise Exception( "Client Message Decrypt Error!" )
	else:
		if server_decrypt:
			if server_decrypt.run(raw_data, raw_size):
				raw_data = server_decrypt.result()
			else:
				dump_message(raw_data, raw_size )
				raise Exception( "Client Message Decrypt Error!" )
	events.append(Packet(svr_ip, svr_port, raw_data, raw_size, captured_raw_size, to_server))
	
raw_size = 0
splitted_data = ''
splitted_size = 0

def recv_packets(header, payload):
	# Read packets -- header contains information about the data from pcap,
	# payload is the actual packet as a string
	global prev_capture_time
	ether = ImpactDecoder.EthDecoder().decode(payload)
	iphdr = ether.child()
	tcphdr = iphdr.child()
	sip = iphdr.get_ip_src()
	sport = tcphdr.get_th_sport()
	dip = iphdr.get_ip_dst()
	dport = tcphdr.get_th_dport()
	total_header_size = ether.get_header_size() + iphdr.get_header_size() + tcphdr.get_header_size()
	for k, v in setting.servers.iteritems():
		svr_ip = v[0]
		svr_port = v[1]
		if (svr_ip == sip and svr_port == sport) or (svr_ip == dip and svr_port == dport):
			if len(payload) > total_header_size:
				svr_addr = svr_ip + str(svr_port)
				to_server = ( svr_ip == dip )
				global local_port
				# capture only first client's packet
				if svr_addr in local_port.keys():
					if to_server:
						if local_port[svr_addr] != sport:
							continue
					else:
						if local_port[svr_addr] != dport:
							continue
				# if first packet, remember local port
				else:
					print "svr_addr: " + svr_addr + " not exist"
					if to_server:
						local_port[svr_addr] = sport
					else:
						local_port[svr_addr] = dport
				print ('\n%s: %s(%s) -> %s(%s) :'
					%(datetime.datetime.now(), sip, sport, dip, dport)),
				print(binascii.hexlify(payload))
				print ('captured %d bytes' % header.getlen()),
				if header.getlen() != header.getcaplen():
					print ('\n** TRUNCATED TO %d BYTES **' % header.getcaplen())
				else:
					print('\n'),
				if to_server:
					events.append(Time(time()-prev_capture_time))
				l = len(payload)
				print(l, total_header_size)
				#print(dir(tcphdr))
				#print(tcphdr.get_bytes())
				print(tcphdr.get_size())
				#print(tcphdr.get_header_size())
				h = tcphdr.get_header_size()
				print( h)
#					payloadLen = iphdr.get_ip_len() - (iphdr.get_ip_hl() * 4) - (tcphdr.get_th_off() * 4)
#					print("payloadLen: ", payloadLen)
				if hasattr(setting, 'get_message_data'):
					raw_data = setting.get_message_data(payload[total_header_size:])
				else:
					raw_data = payload[total_header_size:]				

				captured_raw_size = len(raw_data)
				
				print("len(raw_data) ", len(raw_data))
				
				global raw_size, splitted_data, splitted_size
				print("raw_size ", raw_size)
				# skip empty message
				if setting.is_valid_message_head_size(captured_raw_size):
					# check body of the splitted message
					if splitted_data:
						splitted_data += raw_data
						splitted_size += captured_raw_size
						# check completed
						if splitted_size >= raw_size:
							# composition completed
							raw_data = splitted_data							
							raw_size = splitted_size
							splitted_data = ''
							splitted_size = 0
						else:
							# keep compositing
							return
					else:
						# no splitted message yet
						raw_size = setting.get_message_size(payload[total_header_size:])
						splitted_data = ''
						splitted_size = 0
						
					if raw_size > 0:
						# check head of the splitted message
						if raw_size > captured_raw_size:
							# start compositing
							splitted_data = raw_data
							splitted_size = captured_raw_size
							return
						# check merged message
						elif raw_size < captured_raw_size:
							# append all messages
							while True:
								append_message(raw_data, raw_size, payload, svr_ip, svr_port, captured_raw_size, to_server)
								raw_data, raw_size, next_message_offset = setting.get_next_sibling_message(raw_data, raw_size, captured_raw_size)
								#print "raw_data: " + str(raw_data)
								#print "raw_size: " + str(raw_size)
								captured_raw_size -= next_message_offset
								if captured_raw_size <= 0 or raw_size <= 0:
									break
						# sigle message
						else:
							append_message(raw_data, raw_size, payload, svr_ip, svr_port, captured_raw_size, to_server)
							
						if to_server:
							prev_capture_time = time()
			return
			
def dump_message( message, pl ):
	print binascii.hexlify( message )
	print "\n"

def align_data(data):
	str = ""
	str += '# '
	d = binascii.hexlify(data)
	s = 0
	while d:
		str += '%02s ' % d[:2]
		s += 1
		if s % 32 == 0:
			str += '\n'
			str += '# '
			s = 0
		elif s % 8 == 0:
			str += ' '
		d = d[2:]
	return str

def align_char(data):
	str = ""
	str += '# '
	d = binascii.hexlify(data)
	s = 0
	while d:
		c = int( d[:2], 16 )
		if c < 32 or c > 126:
			c = '.'
		else:
			c = chr(c)
		str += c
		s += 1
		if s % 32 == 0:
			str += '\n'
			str += '# '
			s = 0
		elif s % 8 == 0:
			str += ' '
		d = d[2:]
	return str

def get_message_name( data ):
	name = None
	if hasattr(setting, 'get_message_decorator'):
		decorator = setting.get_message_decorator(data)
		if decorator:
			name = decorator[0]
	if not name:
		m = md5()
		m.update(data)
		name = binascii.hexlify( m.digest() )
	return name

def dump_message_class(f, s, send):
	if send:
		con_str = "dst"
	else:
		con_str = "src"
	message_name = get_message_name(s.data)
	
	if hasattr(setting, 'get_message_decorator'):
		decorator = setting.get_message_decorator(s.data)
		if decorator:
			message_name = decorator[0]
	data = s.data[:s.length]
	f.write( """
%s
%s
class Message_%s(Message):
	length = %s
	captured_data_left = %s
	%s = svr["%s"]
	data = create_string_buffer(a2b_base64('''%s'''))
%s
			""" % (align_data(data), align_char(data), message_name, s.length, s.captured_raw_size, con_str, get_connection_name( s.ip, s.port), 
			binascii.b2a_base64(data), get_property_code(decorator)))

def get_property_code(decorator):
	if not decorator:
		return ''
	ld = len(decorator)
	if ld > 1:
		code = ""
		for i in range(1, ld):
			prop = decorator[i]
			assert( len(prop) == 4)
			name = prop[0]
			start = prop[1]
			length = prop[2]
			ptype = prop[3]
			code += """
	def get_%s(self):
		%s
	def set_%s(self, value):
		%s
	%s = property(get_%s, set_%s)
			""" % ( name, get_get_code(start, length, ptype), name, 
			get_set_code(start, length, ptype), name, name, name )
		return code
	return ""

def get_get_code(start, length, ptype):
	if ptype == 'char':	
		return "return str(self.data[%d:%d+%d-1]).split('\\0')[0]" % (start, start, length)
	elif ptype == 'wstring':
		return "return self.data[%d:%d+%d-2].decode('utf-16')" % (start, start, length)
	elif ptype == 'int':
		return "return unpack('i', self.data[%d:%d+%d])[0]" % (start, start, length)
	elif ptype == 'word':
		return "return unpack('H', self.data[%d:%d+%d])[0]" % (start, start, length)
	elif ptype == 'dword':
		return "return unpack('I', self.data[%d:%d+%d])[0]" % (start, start, length)
	elif ptype == 'byte':
		return "return unpack('B', self.data[%d:%d+%d])[0]" % (start, start, length)
	raise Exception("Undefine type '%s'" % ptype)

def get_set_code(start, length, ptype):
	if ptype == 'char':	
		return """s = value
		assert(len(s)<=%d-1)
		s += '\\0'*(%s-len(s))
		fmt = '%%dc' %% %d
		pack_into(fmt, self.data, %d, *s)""" % (length, length, length, start)
	elif ptype == 'wstring':
		return """s = value.encode('utf-16')[2:]
		assert(len(s)<=%d-2)
		s += '\\0'*(%s-len(s))
		fmt = '%%dc' %% %d
		pack_into(fmt, self.data, %d, *s)""" % (length, length, length, start)
	elif ptype == 'int':
		return "pack_into('i', self.data, %d, value)" % (start)
	elif ptype == 'word':
		return "pack_into('H', self.data, %d, value)" % (start)
	elif ptype == 'dword':
		return "pack_into('I', self.data, %d, value)" % (start)
	elif ptype == 'byte':
		return "pack_into('B', self.data, %d, value)" % (start)
	raise Exception("Undefine type '%s'" % ptype)

def get_connection_name(ip, port):
	for k, v in setting.servers.iteritems():
		if v[0] == ip and v[1] == port:
			return k

def make_pyext( file ):
	return os.path.splitext(file)[0] + ".py"

def flushing(target_file):
	if not target_file:
		target_file = '_captured_.py'
	target_file = os.path.join(os.getcwd(), target_file)
	f = open( target_file, 'wb' )
	f.write( "from binascii import a2b_base64, hexlify\n" )
	f.write( "from struct import unpack_from, pack_into, pack, unpack\n" )
	f.write( "from ctypes import create_string_buffer\n" )
	f.write( "from pyster.connection import Connection as Connection\n")	
	f.write( "from time import time\n")	
	f.write( "import asyncore\n")
	f.write( "\n" )
	f.write( "svr = {}\n" )
	for k, v in setting.servers.iteritems():
		f.write( "svr['%s'] = { 'ip': '%s', 'port': %d }\n" % ( k, v[0], v[1] ) )
	# dump base message object
	f.write( """
# Base Message Object
class Message(object):
	def dump(self):
		print(hexlify(self.data))
	""")
	# dump messages
	for e in events:
		if isinstance( e, Packet):
			p = e
			dump_message_class(f, p, p.send)

	# dump replay
	f.write("\n")
	f.write('con = {}\n')	
	f.write('def replay():\n')
	f.write('	global con\n')
	if len(events):
		inited = {}
		send_delay = 0
		for e in events:
			if isinstance(e, Packet):
				p = e
				message_name = get_message_name(p.data)
				con_name = get_connection_name(p.ip, p.port)
				if not( con_name in inited ):
					ip = p.ip
					port = p.port
					f.write('	con["%s"] = Connection("%s", "%s", %s)\n' % (con_name, con_name, ip, port))
					inited[con_name] = True
				if p.send:
					f.write('	con["%s"].send_with_delay(%f, Message_%s())\n' % (con_name, send_delay, message_name))
					send_delay = 0
				else:
					f.write('	con["%s"].receive("Message_%s")\n' % (con_name, message_name))
			elif isinstance(e, Time):				
				t = e
				send_delay = t.time
	else:
		f.write('	pass')
		
	f.write('\n')
	f.write('if __name__ == "__main__":\n')
	f.write('	replay()\n')
	f.write('	asyncore.loop()\n')

def sniff(target_file):
	import sys
	sys.path.append(os.getcwd())
	global setting
	global client_decrypt, server_decrypt, max_cap_length, filter
	path = os.path.join(os.getcwd(), 'setting.py')
	setting = imp.load_source('setting', path)
	client_decrypt = getattr( setting, 'client_decrypt', None )
	server_decrypt = getattr( setting, 'server_decrypt', None )
	max_cap_length = getattr( setting, 'max_cap_length', 100 )
	filter = getattr( setting, 'filter', 'tcp' )

	device = get_device()
	assert device
	assert len(setting.servers) > 0

	print "Start Sniffing... ( Press Ctrl+C to Stop )"
	try:
		capture(device)
	except KeyboardInterrupt:
		flushing(target_file)

